#include <iostream>
#include <assert.h>
using namespace std;

template <class K, class V>
struct AVLTreeNode{
  AVLTreeNode<K,V> *_left;
  AVLTreeNode<K,V> *_right;
  AVLTreeNode<K,V> *_parent;

  pair<K,V> _kv;
  int _bf;
  AVLTreeNode(const pair<K,V> &kv)
    :_left(nullptr),
    _right(nullptr),
    _parent(nullptr),
    _kv(kv),
    _bf(0)
  {}
};

template <class K, class V>
class AVLTree{
  typedef AVLTreeNode<K,V> Node;
  Node *_root = nullptr;
public:
  bool Insert(const pair<K,V> &kv);
  void Inorder(){
    _Inorder(_root);
    cout << endl;
  }
  int Height(){
    return _Height(_root);
  }
  bool Isbalance(){
    return _Isbalance(_root);
  }
private:
  void RotateL(Node *parent);
  void RotateR(Node *parent);
  void RotateLR(Node *parent);
  void RotateRL(Node *parent);
  void _Inorder(Node *root);
  int _Height(Node *root);
  bool _Isbalance(Node *root);
};


template <class K, class V>
bool AVLTree<K,V>::_Isbalance(Node *root){
  if(root == nullptr) return true;
  int lh = _Height(root->_left);
  int rh = _Height(root->_right);
  int diff = rh-lh;
  if(diff != root->_bf)
  {
    cout << "Key: " << root->_kv.first << "bf: " << root->_bf << " 平衡因子异常" << endl;
    return false;
  }
  return abs(diff) < 2 && _Isbalance(root->_left) && _Isbalance(root->_right);
}

template <class K, class V>
int AVLTree<K,V>::_Height(Node *root){
  if(root == nullptr)
  {
    return 0;
  }
  int lefth = _Height(root->_left);
  int righth = _Height(root->_right);
  return lefth > righth? lefth+1 : righth+1;
}

template <class K, class V>
void AVLTree<K,V>::_Inorder(Node *root){
  if(root == nullptr) return; 
  _Inorder(root->_left);
  cout << root->_kv.first << " ";
  _Inorder(root->_right);
}

template <class K, class V>
bool AVLTree<K,V>::Insert(const pair<K,V> &kv)
{
  if(_root == nullptr)
  {
    _root = new Node(kv);
    return true;
  }

  Node *cur = _root;
  Node *parent = nullptr;
  while(cur != nullptr)
  {
    if(kv.first > cur->_kv.first)
    {
      parent = cur;
      cur = cur->_right;
    }
    else if(kv.first < cur->_kv.first)
    {
      parent  = cur;
      cur = cur->_left;
    }
    else{
      return false;
    }
  }
  cur = new Node(kv);
  if(kv.first > parent->_kv.first)
  {
    parent->_right = cur;
  }
  else{
    parent->_left = cur;
  }
  cur->_parent = parent;
  //控制平衡：1.更新平衡因子
  while(parent!=nullptr)
  {
    if(cur == parent->_right)
    {
      ++parent->_bf;
    }
    else{
      --parent->_bf;
    }

    if(parent->_bf == 0)
    {
      break;  
    }
    else if(abs(parent->_bf) == 1)
    {
      parent = parent->_parent;
      cur = cur->_parent;
    }
    else if(abs(parent->_bf) == 2)
    {
      //说明parent所在的子树已经不平衡了，需要做旋转处理
      if(parent->_bf == 2 && cur->_bf == 1)
      {
        RotateL(parent);
      }
      else if(parent->_bf == -2 && cur->_bf == -1)
      {
        RotateR(parent);
      }
      else if(parent->_bf == -2 && cur->_bf == 1)
      {
        RotateLR(parent);
      }
      else if(parent->_bf == 2 && cur->_bf == -1)
      {
        RotateRL(parent);
      }
      else
      {
        assert(false);
      }
      break;
    }

    else{
      assert(false);
    }
  }
  return true;

}

template <class K, class V>
void AVLTree<K,V>::RotateL(Node *parent){
  Node *subR = parent->_right;
  Node *subRL = subR->_left;
  Node *ppNode = parent->_parent;

  parent->_right = subRL;
  if(subRL != nullptr)
  {
    subRL->_parent = parent;
  }

  subR->_left = parent;
  parent->_parent = subR;
 
  if(_root == parent)
  {
    _root = subR;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subR;
    }
    else{
      ppNode->_right = subR;
    }
  }
  subR->_parent = ppNode;
  subR->_bf = parent->_bf = 0;
}

template <class K, class V>
void AVLTree<K,V>::RotateR(Node *parent){
  Node *subL = parent->_left;
  Node *subLR = subL->_right;
  Node *ppNode = parent->_parent;

  parent->_parent = subL;
  subL->_right = parent;

  parent->_left = subLR;
  if(subLR != nullptr)
  subLR->_parent = parent;

  if(_root == parent)
  {
    _root = subL;
  }
  else{
    if(ppNode->_left == parent)
    {
      ppNode->_left = subL;
    }
    else{
      ppNode->_right = subL;
    }
  }
  subL->_parent = ppNode;
  subL->_bf = parent->_bf = 0;
}

template <class K, class V>
void AVLTree<K,V>::RotateLR(Node *parent){
  Node *subL = parent->_left;
  Node *subLR = subL->_right;
  int bf = subLR->_bf;
  RotateL(subL);
  RotateR(parent);
  subLR->_bf = 0;
  switch(bf)
  {
    case 1:
      subL->_bf = -1;
      parent->_bf = 0;
      break;
    case -1:
      subL->_bf = 0;
      parent->_bf = 1;
      break;
    case 0:
      subL->_bf = 0;
      parent->_bf = 0;
      break;
    default:
      assert(false);
      break;
  }
}

template <class K, class V>
void AVLTree<K,V>::RotateRL(Node *parent){
  Node *subR = parent->_right;
  Node *subRL = subR->_left;
  int bf = subRL->_bf;
  RotateR(subR);
  RotateL(parent);
  subRL->_bf = 0;
  switch(bf)
  {
    case 1:
      subR->_bf = 0;
      parent->_bf = -1;
      break;
    case -1:
      subR->_bf = 1;
      parent->_bf = 0;
      break;
    case 0:
      subR->_bf = 0;
      parent->_bf = 0;
      break;
    default:
      assert(false);
      break;
  }
}
